package com.xialuo.study.agent;

import com.xialuo.study.nodes.LlmNode;
import com.xialuo.study.nodes.ToolNode;
import com.xialuo.study.state.ReactState;
import com.xialuo.study.std.StateSerializer;
import org.bsc.langgraph4j.CompiledGraph;
import org.bsc.langgraph4j.GraphStateException;
import org.bsc.langgraph4j.StateGraph;
import org.bsc.langgraph4j.action.EdgeAction;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.tool.resolution.ToolCallbackResolver;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.bsc.langgraph4j.StateGraph.END;
import static org.bsc.langgraph4j.StateGraph.START;
import static org.bsc.langgraph4j.action.AsyncEdgeAction.edge_async;
import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;

/**
 * 基于LangGraph4j和Spring AI的ReactAgent实现。
 * 该Agent实现了ReAct模式（推理和行动），用于基于LLM的智能体。
 */
public class ReactAgent {

    private String name;
    private final LlmNode llmNode;
    private final ToolNode toolNode;
    private final StateGraph<ReactState> graph;
    private CompiledGraph<ReactState> compiledGraph;
    private int maxIterations = 10;
    private int iterations = 0;

    /**
     * 使用指定参数创建ReactAgent
     *
     * @param name          Agent的名称
     * @param prompt        使用的提示模板
     * @param chatClient    Spring AI的ChatClient
     * @param tools         使用的函数回调列表
     * @param maxIterations 最大迭代次数
     */
    public ReactAgent(String name, String prompt, ChatClient chatClient, List<FunctionCallback> tools, int maxIterations) throws GraphStateException {
        this.name = name;
        this.llmNode = new LlmNode(chatClient, prompt);
        this.toolNode = new ToolNode(tools);
        this.maxIterations = maxIterations;
        this.graph = initGraph();
    }

    /**
     * 使用指定参数创建ReactAgent
     *
     * @param name          Agent的名称
     * @param prompt        使用的提示模板
     * @param chatClient    Spring AI的ChatClient
     * @param resolver      工具回调解析器
     * @param maxIterations 最大迭代次数
     */
    public ReactAgent(String name, String prompt, ChatClient chatClient, ToolCallbackResolver resolver, int maxIterations) throws GraphStateException {
        this.name = name;
        this.llmNode = new LlmNode(chatClient, prompt);
        this.toolNode = new ToolNode(resolver);
        this.maxIterations = maxIterations;
        this.graph = initGraph();
    }

    //复杂图
    /**
     * 构建工作流图，定义节点和转换
     *
     * @return 构建的StateGraph
     * @throws GraphStateException 如果图无法构建
     */
//    public StateGraph<State> build() throws GraphStateException {
//        var routeToNextAgent = (EdgeAction<State>) state -> {
//            log.info("路由到下一个Agent: {}", state);
//            String next = state.getNext();
//            if ("FINISH".equals(next)) {
//                return "end";
//            }
//            return next;
//        };
//
//        return new StateGraph<>(State.SCHEMA, new StateSerializer())
//                .addEdge(START, "supervisor")
//                .addNode("supervisor", node_async(AgentExecutor.this::callSupervisorNode))
//                .addConditionalEdges("supervisor",
//                        edge_async(routeToNextAgent),
//                        Map.of(
//                                "researcher", "researcher",
//                                "coder", "coder",
//                                "datetime_worker", "datetime_worker",
//                                "end", END
//                        )
//                )
//                .addNode("researcher", node_async(AgentExecutor.this::callResearcherNode))
//                .addEdge("researcher", "supervisor")
//                .addNode("coder", node_async(AgentExecutor.this::callCoderNode))
//                .addEdge("coder", "supervisor")
//                .addNode("datetime_worker", node_async(AgentExecutor.this::callDateTimeNode))
//                .addEdge("datetime_worker", "supervisor");
//    }

    /**
     * 初始化Agent的状态图
     *
     * @return 初始化的状态图
     */
    private StateGraph<ReactState> initGraph() throws GraphStateException {
        return new StateGraph<>(ReactState.SCHEMA, new StateSerializer())
                .addNode("agent", this.llmNode)
                .addNode("tool", this.toolNode)
                .addEdge(START, "agent")
                .addConditionalEdges("agent",
                                     edge_async(this::think),
                                     Map.of("continue", "tool",
                                            "end", END
                                     )
                )
                .addEdge("tool", "agent");
    }

    /**
     * 决策函数，用于确定Agent是否应该继续或结束
     *
     * @param state 当前Agent状态
     * @return 如果Agent应该继续，则返回"continue"，否则返回"end"
     */
    private String think(ReactState state) {
        if (iterations > maxIterations) {
            return "end";
        }

        List<Message> messages = state.getMessages();
        if (messages == null || messages.isEmpty()) {
            return "end";
        }

        Message lastMessage = messages.get(messages.size() - 1);
        if (lastMessage instanceof AssistantMessage) {
            AssistantMessage assistantMessage = (AssistantMessage) lastMessage;
            if (assistantMessage.hasToolCalls()) {
                return "continue";
            }
        }

        return "end";
    }

    /**
     * 编译图并返回
     *
     * @return 编译后的图
     */
    public CompiledGraph<ReactState> getCompiledGraph() throws GraphStateException {
        if (this.compiledGraph == null) {
            this.compiledGraph = this.graph.compile();
        }
        return this.compiledGraph;
    }

    /**
     * 使用给定输入调用Agent
     *
     * @param input 输入消息
     * @return Agent的最终响应
     */
    public String invoke(String input) throws GraphStateException {
        CompiledGraph<ReactState> compiledGraph = getCompiledGraph();

        // 创建包含用户消息的初始状态
        Map<String, Object> initialState = new HashMap<>();
        List<Message> messages = new ArrayList<>();
        messages.add(new UserMessage(input));
        initialState.put("messages", messages);

        // 调用图
        ReactState finalState = compiledGraph
                .invoke(initialState)
                .get();

        // 提取最终响应
        List<Message> resultMessages = finalState.getMessages();
        AssistantMessage lastMessage = (AssistantMessage) resultMessages.get(resultMessages.size() - 1);
        return lastMessage.getText();
    }

    /**
     * 创建新的Builder实例
     *
     * @return 新的Builder实例
     */
    public static ReactAgentBuilder builder() {
        return new ReactAgentBuilder();
    }
}
